Skip to content

[SPARK-61152][CONNECT][PYTHON] Raise INVALID_RETURN_TYPE_FOR_PANDAS_UDF for a non-struct return type of mapInPandas, applyInPandas and the other map functions, in the Python client as PySpark does - #341

Closed
HyukjinKwon wants to merge 1 commit into
apache:masterfrom
HyukjinKwon:SPARK-61152-map-functions-struct-return-type
Closed

HyukjinKwon wants to merge 1 commit into
apache:masterfrom
HyukjinKwon:SPARK-61152-map-functions-struct-return-type

Conversation

@HyukjinKwon

Copy link
Copy Markdown
Member

What changes were proposed in this pull request?

SPARK-61152: Raise INVALID_RETURN_TYPE_FOR_PANDAS_UDF for a non-struct return type of mapInPandas, applyInPandas and the other map functions, in the Python client as PySpark does.

PySpark builds the UDF of mapInPandas, mapInArrow, applyInPandas, applyInArrow, applyInPandasWithState and the cogrouped applyInPandas and applyInArrow through UserDefinedFunction, which checks the return type for the eval type. A return type that is not a struct raises PySparkTypeError INVALID_RETURN_TYPE_FOR_PANDAS_UDF, or INVALID_RETURN_TYPE_FOR_ARROW_UDF for the Arrow functions, at the call, and the plan is never sent.

The drop-in already has this check verbatim in pyspark/sql/udf.py. build_map_udf now calls UserDefinedFunction._check_return_type with the resolved return type and the eval type, right after the schema is resolved, as the official client does. The check is a no-op for the transform-with-state eval types, which build with build_map_udf too.

Why are the changes needed?

df.mapInPandas(f, "int"), df.mapInPandas(f, IntegerType()), df.mapInArrow(f, "long"), df.groupBy("id").applyInPandas(g, "int") and df.groupBy("id").applyInArrow(g, LongType()) were sent to the server. The server failed with UnknownException (ClassCastException), or IllegalArgumentException: The returnType of the udf must be a struct. PySpark raises PySparkTypeError INVALID_RETURN_TYPE_FOR_PANDAS_UDF (for example Pandas UDF should return StructType for SQL_MAP_PANDAS_ITER_UDF or SQL_MAP_ARROW_ITER_UDF, got IntegerType().) on the client.

Does this PR introduce any user-facing change?

Yes. A map or group function with a non-struct return type raises the PySpark PySparkTypeError on the client when the call builds the plan, with the same error class, message and return_type as PySpark, instead of a server error.

How was this patch tested?

  • test_map_and_group_functions_reject_a_non_struct_return_type_as_pyspark (offline, mock server): the nine call forms (mapInPandas with a type and a DDL string, mapInArrow, applyInPandas with a DDL string and an ArrayType, applyInArrow, the two cogrouped applies, applyInPandasWithState) raise the expected condition and return_type at the call, and no request is sent. The same calls give the same condition, message and return type with official PySpark 4.3.0.
  • test_map_functions_reject_a_non_struct_return_type_before_the_server (live): mapInPandas and applyInArrow with a non-struct type raise on the client, and a struct schema still runs.
  • The lib tests of both crates (590 passed), the build of all tests, also with wasm-udf, and the offline Python suites (2202 passed, 8 skipped), and the live tests of the map, grouped, cogrouped and stateful functions (22 passed) against a local Spark 4.3.0 server; rustfmt passes, and clippy reports no diagnostics in the changed lines.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Opus 5.5

This pull request and its description were written by Isaac.

…DF for a non-struct return type of mapInPandas, applyInPandas and the other map functions, in the Python client as PySpark does

mapInPandas, mapInArrow, applyInPandas, applyInArrow, applyInPandasWithState and the cogrouped applyInPandas and applyInArrow sent a plan whose UDF returns a type that is not a struct, and the server failed with a ClassCastException or an IllegalArgumentException. As PySpark, the UDF build now checks its return type with UserDefinedFunction._check_return_type, which raises INVALID_RETURN_TYPE_FOR_PANDAS_UDF or INVALID_RETURN_TYPE_FOR_ARROW_UDF on the client, at the call.

Co-authored-by: Isaac <no-reply@databricks.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant